import os
import json
from run import run_and_return_test
import pandas as pd
from submission import write_ensemble_submission
import numpy as np
import gc
from config_colab import *

prefix = "/content/drive/MyDrive/lab-result-1/groups"

# fold_num = "0"
test_probs = None

# for fold_num in ["0","1","2","3","4"]:
#     config["architecture"] ="/content/drive/MyDrive/ccks2021-wws/data/test_data_roformer_char_fulltoken"
#     config["train_name"] = config["architecture"].split("/")[-1]+"_sapmle_weight_fold_"+fold_num

#     config["inputs_fold"] = f"/content/drive/MyDrive/ccks2021-wws/lab/data/5fold/fold_"+fold_num

#     path = os.path.join(prefix, config["group_name"], config["train_name"])
#     config["output_path"] = os.path.join(path, "outputs")
#     config_path = os.path.join(path, "config.json")
#     if not os.path.exists(path):
#         os.makedirs(config["output_path"])
#     json.dump(config, open(os.path.join(path, "config.json"), mode="w"))

#     if test_probs is None:
#         test_probs = run_and_return_test(config_path)
#     else:
#         test_probs = test_probs+ run_and_return_test(config_path)

#     np.save("/content/drive/MyDrive/ccks2021-wws/test_probs.npy",test_probs)
#     write_ensemble_submission(test_probs)

#     gc.collect()
# config["train_fold"] =architecture
for group in ["a1","a2","a3","a4","a5","a6"]:
    submission_path = f"/content/drive/MyDrive/ccks2021-wws/submission-{group}-times.txt"

    for architecture in [
        # "/content/drive/MyDrive/ccks2021-wws/data/ext_data_roformer_char_fulltoken",
        "/content/drive/MyDrive/ccks2021-wws/data/test_data_roformer_char_fulltoken"
        # "/content/drive/MyDrive/ccks2021-wws/data/output_roformer_chinese_base",
        # # "/content/drive/MyDrive/ccks2021-wws/data/output_nghuyong-ernie-1.0",
        # "/content/drive/MyDrive/ccks2021-wws/data/output_hfl-chinese-roberta-wwm-ext",
        # "/content/drive/MyDrive/ccks2021-wws/data/output_mlm_hfl-chinese-bert-wwm"
    ]:

        for fold_num in ["0","1","2","3","4"]:
            config["architecture"] =architecture
            config["group_name"] = group
            config["train_name"] = architecture.split("/")[-1]+"_fold_"+fold_num

            config["inputs_fold"] = f"/content/drive/MyDrive/ccks2021-wws/lab/data/5fold/fold_"+fold_num

            path = os.path.join(prefix, config["group_name"], config["train_name"])
            config["output_path"] = os.path.join(path, "outputs")
            config_path = os.path.join(path, "config.json")
            if not os.path.exists(path):
                os.makedirs(config["output_path"])
            json.dump(config, open(os.path.join(path, "config.json"), mode="w"))

            if test_probs is None:
                test_probs = run_and_return_test(config_path)
            else:
                test_probs = test_probs+ run_and_return_test(config_path)

            np.save("/content/drive/MyDrive/ccks2021-wws/test_probs-ext.npy",test_probs)
            write_ensemble_submission(test_probs,submission_path=submission_path)

            gc.collect()

    Config.fake_label_test_path = submission_path


    
    

